{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "51466c8d-8ce4-4b3d-be4e-18fdbeda5f53",
   "metadata": {},
   "source": [
    "# How to use Postgres checkpointer for persistence\n",
    "\n",
    "<div class=\"admonition tip\">\n",
    "    <p class=\"admonition-title\">Prerequisites</p>\n",
    "    <p>\n",
    "        This guide assumes familiarity with the following:\n",
    "        <ul>\n",
    "            <li>\n",
    "                <a href=\"https://langchain-ai.github.io/langgraphjs/concepts/persistence/\">\n",
    "                    Persistence\n",
    "                </a>\n",
    "            </li>       \n",
    "            <li>\n",
    "                <a href=\"https://www.postgresql.org/about/\">\n",
    "                    Postgresql\n",
    "                </a>\n",
    "            </li>        \n",
    "        </ul>\n",
    "    </p>\n",
    "</div> \n",
    "\n",
    "When creating LangGraph agents, you can set them up so that they persist their state across executions. This allows you to do things like interact with an agent multiple times and have it remember previous interactions.\n",
    "\n",
    "This how-to guide shows how to use Postgres as the backend for persisting checkpoint state using the [`@langchain/langgraph-checkpoint-postgres`](https://github.com/langchain-ai/langgraphjs/tree/main/libs/checkpoint-postgres) library and the [`PostgresSaver`](https://langchain-ai.github.io/langgraphjs/reference/classes/checkpoint_postgres.PostgresSaver.html) class.\n",
    "\n",
    "For demonstration purposes we will add persistence to the [pre-built create react agent](https://langchain-ai.github.io/langgraphjs/reference/functions/langgraph_prebuilt.createReactAgent.html). \n",
    "\n",
    "In general, you can add a checkpointer to any custom graph that you build like this:\n",
    "\n",
    "```ts\n",
    "import { StateGraph } from \"@langchain/langgraph\";\n",
    "import { PostgresSaver } from \"@langchain/langgraph-checkpoint-postgres\";\n",
    "\n",
    "const builder = new StateGraph(...);\n",
    "\n",
    "// ... define the graph\n",
    "\n",
    "const checkpointer = PostgresSaver.fromConnString(...); // postgres checkpointer (see examples below)\n",
    "\n",
    "const graph = builder.compile({ checkpointer });\n",
    "...\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "456fa19c-93a5-4750-a410-f2d810b964ad",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "You will need access to a Postgres instance. This guide will also use OpenAI, so you will need an OpenAI API key.\n",
    "\n",
    "First, install the required packages:\n",
    "\n",
    "```bash\n",
    "npm install @langchain/langgraph @langchain/core @langchain/langgraph-checkpoint-postgres\n",
    "```\n",
    "\n",
    "Then, set your OpenAI API key as `process.env.OPENAI_API_KEY`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b394e26c",
   "metadata": {},
   "source": [
    "<div class=\"admonition tip\">\n",
    "    <p class=\"admonition-title\">Set up <a href=\"https://smith.langchain.com\">LangSmith</a> for LangGraph development</p>\n",
    "    <p style=\"padding-top: 5px;\">\n",
    "        Sign up for LangSmith to quickly spot issues and improve the performance of your LangGraph projects. LangSmith lets you use trace data to debug, test, and monitor your LLM apps built with LangGraph — read more about how to get started <a href=\"https://docs.smith.langchain.com\">here</a>. \n",
    "    </p>\n",
    "</div>    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e26b3204-cca2-414c-800e-7e09032445ae",
   "metadata": {},
   "source": [
    "## Define model and tools for the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e5213193-5a7d-43e7-aeba-fe732bb1cd7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import { z } from \"zod\";\n",
    "import { tool } from \"@langchain/core/tools\";\n",
    "\n",
    "const getWeather = tool(async (input: { city: \"sf\" | \"nyc\" }) => {\n",
    "  if (input.city === \"nyc\") {\n",
    "    return \"It might be cloudy in nyc\";\n",
    "  } else if (input.city === \"sf\") {\n",
    "    return \"It's always sunny in sf\";\n",
    "  } else {\n",
    "    throw new Error(\"Unknown city\");\n",
    "  }\n",
    "}, {\n",
    "  name: \"get_weather\",\n",
    "  description: \"Use this to get weather information.\",\n",
    "  schema: z.object({\n",
    "    city: z.enum([\"sf\", \"nyc\"])\n",
    "  }),\n",
    "});"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9342c62-dbb4-40f6-9271-7393f1ca48c4",
   "metadata": {},
   "source": [
    "## With a connection pool\n",
    "\n",
    "Under the hood, `PostgresSaver` uses the [`node-postgres`](https://www.npmjs.com/package/pg) (`pg`) package to connect to your Postgres instance. You can pass in a [connection pool](https://node-postgres.com/apis/pool) that you've instantiated like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2b9d13b1-9d72-48a0-b63a-adc062c06c29",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "  messages: [\n",
      "    HumanMessage {\n",
      "      \"id\": \"ac832b73-242d-4d0b-80d7-5d06a908787e\",\n",
      "      \"content\": \"what's the weather in sf\",\n",
      "      \"additional_kwargs\": {},\n",
      "      \"response_metadata\": {}\n",
      "    },\n",
      "    AIMessage {\n",
      "      \"id\": \"chatcmpl-AGC3tgRXInGLo0qzrD5u3gNqNOegf\",\n",
      "      \"content\": \"\",\n",
      "      \"additional_kwargs\": {\n",
      "        \"tool_calls\": [\n",
      "          {\n",
      "            \"id\": \"call_I2Ceef2LoxjeaR9m8ZkY7U1R\",\n",
      "            \"type\": \"function\",\n",
      "            \"function\": \"[Object]\"\n",
      "          }\n",
      "        ]\n",
      "      },\n",
      "      \"response_metadata\": {\n",
      "        \"tokenUsage\": {\n",
      "          \"completionTokens\": 14,\n",
      "          \"promptTokens\": 57,\n",
      "          \"totalTokens\": 71\n",
      "        },\n",
      "        \"finish_reason\": \"tool_calls\",\n",
      "        \"system_fingerprint\": \"fp_f85bea6784\"\n",
      "      },\n",
      "      \"tool_calls\": [\n",
      "        {\n",
      "          \"name\": \"get_weather\",\n",
      "          \"args\": {\n",
      "            \"city\": \"sf\"\n",
      "          },\n",
      "          \"type\": \"tool_call\",\n",
      "          \"id\": \"call_I2Ceef2LoxjeaR9m8ZkY7U1R\"\n",
      "        }\n",
      "      ],\n",
      "      \"invalid_tool_calls\": [],\n",
      "      \"usage_metadata\": {\n",
      "        \"input_tokens\": 57,\n",
      "        \"output_tokens\": 14,\n",
      "        \"total_tokens\": 71\n",
      "      }\n",
      "    },\n",
      "    ToolMessage {\n",
      "      \"id\": \"6533d271-6126-40af-b5d0-23a484853a97\",\n",
      "      \"content\": \"It's always sunny in sf\",\n",
      "      \"name\": \"get_weather\",\n",
      "      \"additional_kwargs\": {},\n",
      "      \"response_metadata\": {},\n",
      "      \"tool_call_id\": \"call_I2Ceef2LoxjeaR9m8ZkY7U1R\"\n",
      "    },\n",
      "    AIMessage {\n",
      "      \"id\": \"chatcmpl-AGC3ttvB69pQu0atw0lUzTpNePlPn\",\n",
      "      \"content\": \"The weather in San Francisco (SF) is always sunny!\",\n",
      "      \"additional_kwargs\": {},\n",
      "      \"response_metadata\": {\n",
      "        \"tokenUsage\": {\n",
      "          \"completionTokens\": 13,\n",
      "          \"promptTokens\": 84,\n",
      "          \"totalTokens\": 97\n",
      "        },\n",
      "        \"finish_reason\": \"stop\",\n",
      "        \"system_fingerprint\": \"fp_f85bea6784\"\n",
      "      },\n",
      "      \"tool_calls\": [],\n",
      "      \"invalid_tool_calls\": [],\n",
      "      \"usage_metadata\": {\n",
      "        \"input_tokens\": 84,\n",
      "        \"output_tokens\": 13,\n",
      "        \"total_tokens\": 97\n",
      "      }\n",
      "    }\n",
      "  ]\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "import { ChatOpenAI } from \"@langchain/openai\";\n",
    "import { PostgresSaver } from \"@langchain/langgraph-checkpoint-postgres\";\n",
    "import { createReactAgent } from \"@langchain/langgraph/prebuilt\";\n",
    "\n",
    "import pg from \"pg\";\n",
    "\n",
    "const { Pool } = pg;\n",
    "\n",
    "const pool = new Pool({\n",
    "  connectionString: \"postgresql://user:password@localhost:5434/testdb\"\n",
    "});\n",
    "\n",
    "const checkpointer = new PostgresSaver(pool);\n",
    "\n",
    "// NOTE: you need to call .setup() the first time you're using your checkpointer\n",
    "\n",
    "await checkpointer.setup();\n",
    "\n",
    "const graph = createReactAgent({\n",
    "  tools: [getWeather],\n",
    "  llm: new ChatOpenAI({\n",
    "    model: \"gpt-4o-mini\",\n",
    "  }),\n",
    "  checkpointSaver: checkpointer,\n",
    "});\n",
    "const config = { configurable: { thread_id: \"1\" } };\n",
    "\n",
    "await graph.invoke({\n",
    "  messages: [{\n",
    "    role: \"user\",\n",
    "    content: \"what's the weather in sf\"\n",
    "  }],\n",
    "}, config);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1f7a2f6b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "  v: 1,\n",
      "  id: '1ef85bc6-bd28-67c1-8003-5cb7dab561b0',\n",
      "  ts: '2024-10-08T21:29:38.109Z',\n",
      "  pending_sends: [],\n",
      "  versions_seen: {\n",
      "    agent: { tools: 4, '__start__:agent': 2 },\n",
      "    tools: { 'branch:agent:shouldContinue:tools': 3 },\n",
      "    __input__: {},\n",
      "    __start__: { __start__: 1 }\n",
      "  },\n",
      "  channel_versions: {\n",
      "    agent: 5,\n",
      "    tools: 5,\n",
      "    messages: 5,\n",
      "    __start__: 2,\n",
      "    '__start__:agent': 3,\n",
      "    'branch:agent:shouldContinue:tools': 4\n",
      "  },\n",
      "  channel_values: {\n",
      "    agent: 'agent',\n",
      "    messages: [\n",
      "      HumanMessage {\n",
      "        \"id\": \"ac832b73-242d-4d0b-80d7-5d06a908787e\",\n",
      "        \"content\": \"what's the weather in sf\",\n",
      "        \"additional_kwargs\": {},\n",
      "        \"response_metadata\": {}\n",
      "      },\n",
      "      AIMessage {\n",
      "        \"id\": \"chatcmpl-AGC3tgRXInGLo0qzrD5u3gNqNOegf\",\n",
      "        \"content\": \"\",\n",
      "        \"additional_kwargs\": {\n",
      "          \"tool_calls\": [\n",
      "            {\n",
      "              \"id\": \"call_I2Ceef2LoxjeaR9m8ZkY7U1R\",\n",
      "              \"type\": \"function\",\n",
      "              \"function\": \"[Object]\"\n",
      "            }\n",
      "          ]\n",
      "        },\n",
      "        \"response_metadata\": {\n",
      "          \"tokenUsage\": {\n",
      "            \"completionTokens\": 14,\n",
      "            \"promptTokens\": 57,\n",
      "            \"totalTokens\": 71\n",
      "          },\n",
      "          \"finish_reason\": \"tool_calls\",\n",
      "          \"system_fingerprint\": \"fp_f85bea6784\"\n",
      "        },\n",
      "        \"tool_calls\": [\n",
      "          {\n",
      "            \"name\": \"get_weather\",\n",
      "            \"args\": {\n",
      "              \"city\": \"sf\"\n",
      "            },\n",
      "            \"type\": \"tool_call\",\n",
      "            \"id\": \"call_I2Ceef2LoxjeaR9m8ZkY7U1R\"\n",
      "          }\n",
      "        ],\n",
      "        \"invalid_tool_calls\": []\n",
      "      },\n",
      "      ToolMessage {\n",
      "        \"id\": \"6533d271-6126-40af-b5d0-23a484853a97\",\n",
      "        \"content\": \"It's always sunny in sf\",\n",
      "        \"name\": \"get_weather\",\n",
      "        \"additional_kwargs\": {},\n",
      "        \"response_metadata\": {},\n",
      "        \"tool_call_id\": \"call_I2Ceef2LoxjeaR9m8ZkY7U1R\"\n",
      "      },\n",
      "      AIMessage {\n",
      "        \"id\": \"chatcmpl-AGC3ttvB69pQu0atw0lUzTpNePlPn\",\n",
      "        \"content\": \"The weather in San Francisco (SF) is always sunny!\",\n",
      "        \"additional_kwargs\": {},\n",
      "        \"response_metadata\": {\n",
      "          \"tokenUsage\": {\n",
      "            \"completionTokens\": 13,\n",
      "            \"promptTokens\": 84,\n",
      "            \"totalTokens\": 97\n",
      "          },\n",
      "          \"finish_reason\": \"stop\",\n",
      "          \"system_fingerprint\": \"fp_f85bea6784\"\n",
      "        },\n",
      "        \"tool_calls\": [],\n",
      "        \"invalid_tool_calls\": []\n",
      "      }\n",
      "    ]\n",
      "  }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "await checkpointer.get(config);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49fb52fd-af31-4603-889d-66d783244bce",
   "metadata": {},
   "source": [
    "### With a connection string\n",
    "\n",
    "You can also create a pool internally by passing a connection string to the `.fromConnString` static method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5fe54e79-9eaf-44e2-b2d9-1e0284b984d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "  messages: [\n",
      "    HumanMessage {\n",
      "      \"id\": \"c17b65af-6ac5-411e-ab5c-8003dc53755d\",\n",
      "      \"content\": \"what's the weather in sf\",\n",
      "      \"additional_kwargs\": {},\n",
      "      \"response_metadata\": {}\n",
      "    },\n",
      "    AIMessage {\n",
      "      \"id\": \"chatcmpl-AGC6n8XO05i1Z7f4GnOqpayLPxgoF\",\n",
      "      \"content\": \"\",\n",
      "      \"additional_kwargs\": {\n",
      "        \"tool_calls\": [\n",
      "          {\n",
      "            \"id\": \"call_n9QCrJ4QbmgFkr5fHEsQHCCO\",\n",
      "            \"type\": \"function\",\n",
      "            \"function\": \"[Object]\"\n",
      "          }\n",
      "        ]\n",
      "      },\n",
      "      \"response_metadata\": {\n",
      "        \"tokenUsage\": {\n",
      "          \"completionTokens\": 14,\n",
      "          \"promptTokens\": 57,\n",
      "          \"totalTokens\": 71\n",
      "        },\n",
      "        \"finish_reason\": \"tool_calls\",\n",
      "        \"system_fingerprint\": \"fp_f85bea6784\"\n",
      "      },\n",
      "      \"tool_calls\": [\n",
      "        {\n",
      "          \"name\": \"get_weather\",\n",
      "          \"args\": {\n",
      "            \"city\": \"sf\"\n",
      "          },\n",
      "          \"type\": \"tool_call\",\n",
      "          \"id\": \"call_n9QCrJ4QbmgFkr5fHEsQHCCO\"\n",
      "        }\n",
      "      ],\n",
      "      \"invalid_tool_calls\": [],\n",
      "      \"usage_metadata\": {\n",
      "        \"input_tokens\": 57,\n",
      "        \"output_tokens\": 14,\n",
      "        \"total_tokens\": 71\n",
      "      }\n",
      "    },\n",
      "    ToolMessage {\n",
      "      \"id\": \"779c26b0-6b75-454e-98ef-ecca79e50e8c\",\n",
      "      \"content\": \"It's always sunny in sf\",\n",
      "      \"name\": \"get_weather\",\n",
      "      \"additional_kwargs\": {},\n",
      "      \"response_metadata\": {},\n",
      "      \"tool_call_id\": \"call_n9QCrJ4QbmgFkr5fHEsQHCCO\"\n",
      "    },\n",
      "    AIMessage {\n",
      "      \"id\": \"chatcmpl-AGC6ngqEV0EBZbPwHf2JgTw0n16D8\",\n",
      "      \"content\": \"The weather in San Francisco (SF) is described as always sunny.\",\n",
      "      \"additional_kwargs\": {},\n",
      "      \"response_metadata\": {\n",
      "        \"tokenUsage\": {\n",
      "          \"completionTokens\": 15,\n",
      "          \"promptTokens\": 84,\n",
      "          \"totalTokens\": 99\n",
      "        },\n",
      "        \"finish_reason\": \"stop\",\n",
      "        \"system_fingerprint\": \"fp_74ba47b4ac\"\n",
      "      },\n",
      "      \"tool_calls\": [],\n",
      "      \"invalid_tool_calls\": [],\n",
      "      \"usage_metadata\": {\n",
      "        \"input_tokens\": 84,\n",
      "        \"output_tokens\": 15,\n",
      "        \"total_tokens\": 99\n",
      "      }\n",
      "    }\n",
      "  ]\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "const checkpointerFromConnString = PostgresSaver.fromConnString(\n",
    "  \"postgresql://user:password@localhost:5434/testdb\"\n",
    ");\n",
    "\n",
    "const graph2 = createReactAgent({\n",
    "  tools: [getWeather],\n",
    "  llm: new ChatOpenAI({\n",
    "    model: \"gpt-4o-mini\",\n",
    "  }),\n",
    "  checkpointSaver: checkpointerFromConnString,\n",
    "});\n",
    "const config2 = { configurable: { thread_id: \"2\" } };\n",
    "\n",
    "await graph2.invoke({\n",
    "  messages: [{\n",
    "    role: \"user\",\n",
    "    content: \"what's the weather in sf\"\n",
    "  }],\n",
    "}, config2);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b2ce743b-5896-443b-9ec0-a655b065895c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "  v: 1,\n",
      "  id: '1ef85bcd-71b9-6671-8003-6e734c8e9679',\n",
      "  ts: '2024-10-08T21:32:38.103Z',\n",
      "  pending_sends: [],\n",
      "  versions_seen: {\n",
      "    agent: { tools: 4, '__start__:agent': 2 },\n",
      "    tools: { 'branch:agent:shouldContinue:tools': 3 },\n",
      "    __input__: {},\n",
      "    __start__: { __start__: 1 }\n",
      "  },\n",
      "  channel_versions: {\n",
      "    agent: 5,\n",
      "    tools: 5,\n",
      "    messages: 5,\n",
      "    __start__: 2,\n",
      "    '__start__:agent': 3,\n",
      "    'branch:agent:shouldContinue:tools': 4\n",
      "  },\n",
      "  channel_values: {\n",
      "    agent: 'agent',\n",
      "    messages: [\n",
      "      HumanMessage {\n",
      "        \"id\": \"c17b65af-6ac5-411e-ab5c-8003dc53755d\",\n",
      "        \"content\": \"what's the weather in sf\",\n",
      "        \"additional_kwargs\": {},\n",
      "        \"response_metadata\": {}\n",
      "      },\n",
      "      AIMessage {\n",
      "        \"id\": \"chatcmpl-AGC6n8XO05i1Z7f4GnOqpayLPxgoF\",\n",
      "        \"content\": \"\",\n",
      "        \"additional_kwargs\": {\n",
      "          \"tool_calls\": [\n",
      "            {\n",
      "              \"id\": \"call_n9QCrJ4QbmgFkr5fHEsQHCCO\",\n",
      "              \"type\": \"function\",\n",
      "              \"function\": \"[Object]\"\n",
      "            }\n",
      "          ]\n",
      "        },\n",
      "        \"response_metadata\": {\n",
      "          \"tokenUsage\": {\n",
      "            \"completionTokens\": 14,\n",
      "            \"promptTokens\": 57,\n",
      "            \"totalTokens\": 71\n",
      "          },\n",
      "          \"finish_reason\": \"tool_calls\",\n",
      "          \"system_fingerprint\": \"fp_f85bea6784\"\n",
      "        },\n",
      "        \"tool_calls\": [\n",
      "          {\n",
      "            \"name\": \"get_weather\",\n",
      "            \"args\": {\n",
      "              \"city\": \"sf\"\n",
      "            },\n",
      "            \"type\": \"tool_call\",\n",
      "            \"id\": \"call_n9QCrJ4QbmgFkr5fHEsQHCCO\"\n",
      "          }\n",
      "        ],\n",
      "        \"invalid_tool_calls\": []\n",
      "      },\n",
      "      ToolMessage {\n",
      "        \"id\": \"779c26b0-6b75-454e-98ef-ecca79e50e8c\",\n",
      "        \"content\": \"It's always sunny in sf\",\n",
      "        \"name\": \"get_weather\",\n",
      "        \"additional_kwargs\": {},\n",
      "        \"response_metadata\": {},\n",
      "        \"tool_call_id\": \"call_n9QCrJ4QbmgFkr5fHEsQHCCO\"\n",
      "      },\n",
      "      AIMessage {\n",
      "        \"id\": \"chatcmpl-AGC6ngqEV0EBZbPwHf2JgTw0n16D8\",\n",
      "        \"content\": \"The weather in San Francisco (SF) is described as always sunny.\",\n",
      "        \"additional_kwargs\": {},\n",
      "        \"response_metadata\": {\n",
      "          \"tokenUsage\": {\n",
      "            \"completionTokens\": 15,\n",
      "            \"promptTokens\": 84,\n",
      "            \"totalTokens\": 99\n",
      "          },\n",
      "          \"finish_reason\": \"stop\",\n",
      "          \"system_fingerprint\": \"fp_74ba47b4ac\"\n",
      "        },\n",
      "        \"tool_calls\": [],\n",
      "        \"invalid_tool_calls\": []\n",
      "      }\n",
      "    ]\n",
      "  }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "await checkpointerFromConnString.get(config2);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d58ea01",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "TypeScript",
   "language": "typescript",
   "name": "tslab"
  },
  "language_info": {
   "codemirror_mode": {
    "mode": "typescript",
    "name": "javascript",
    "typescript": true
   },
   "file_extension": ".ts",
   "mimetype": "text/typescript",
   "name": "typescript",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
